import os

import dragon_ai
from gxl_ai_utils.utils import utils_file

def do_fenfa():
    # 讲副语言数据 s2t都转成s2t thinker的任务模式。
    dict_info = utils_file.load_dict_from_yaml("./data_s2t.yaml")
    no_thinker_tag = "<S2TCHAT>"
    thinker_tag = '<S2TCHAT> <THINKER>'
    for key, info_dict in dict_info.items():
        dragon_ai.print_info(f'Processing {key}...')
        combine_list_path = info_dict['combine_path']
        tar_root_dir_path = combine_list_path.replace('combine_list.txt', 'combines_tar_root.txt')
        tar_root = dragon_ai.load_first_row_from_file_clean(tar_root_dir_path)
        now_dir = os.path.dirname(combine_list_path)
        data_list_path = utils_file.join_path(now_dir, 'data.list')
        dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, os.path.dirname(combine_list_path),
                                                          task_tag=thinker_tag, tar_dir=tar_root)
        no_think_dir = now_dir.replace('s2t_thinking', 's2t_no_thinking')
        dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, no_think_dir, task_tag=no_thinker_tag,
                                                          tar_dir=tar_root)


def get_old_data_shards():
    """"""
    time_now = dragon_ai.do_get_now_time()
    combines_path = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_thinking/doubao/data_old/combines_list.txt"
    output_tar_dir = "/home/A02_tmpdata3/osum_s2s/gender_old"
    output_wav_dir = f'{output_tar_dir}/wav'
    dragon_ai.makedirs(output_wav_dir)
    old_tar_path = "/home/A02_tmpdata3/osum_s2s/gender"
    tar_path_all_list_tmp = dragon_ai.load_list_from_file_clean(combines_path)
    tar_path_all_list = []
    for tar_path in tar_path_all_list_tmp:
        file_name_pure = dragon_ai.do_get_file_pure_name_from_path(tar_path)
        new_tar_path = os.path.join(old_tar_path, file_name_pure + ".tar")
        tar_path_all_list.append(new_tar_path)
    fake_path = dragon_ai.do_get_fake_file_path()
    dragon_ai.write_list_to_file(tar_path_all_list, fake_path)
    test_data_loader = dragon_ai.do_get_data_loader_for_read(fake_path, batch_size=50, num_workers=10)

    old_data_list = dragon_ai.load_dict_list_from_jsonl("/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_thinking/doubao/data_old/data.list")
    keys_set = {itm['key'] for itm in old_data_list}

    data_list = []
    for i, batch in enumerate(test_data_loader):
        sorted_key = batch['keys']
        sorted_txt = batch['labels']
        sorted_wav = batch['wavs']
        sorted_extra_dicts = batch['extra_dicts']
        sorted_file_path = batch['tar_files']
        j = 0
        for key, txt, wav ,extra_dict,file_path in zip(sorted_key, sorted_txt, sorted_wav,sorted_extra_dicts,sorted_file_path):
            if key not in keys_set:
                continue
            j += 1
            # print_info(f'j: {j}, key: {key}, txt: {txt}')
            wav_path = os.path.join(output_wav_dir, f'{key}.wav')
            file_path_key = dragon_ai.do_get_file_pure_name_from_path(file_path)
            item_dict = {'key': key, 'txt': txt, 'wav': wav_path, 'extra_dict': extra_dict, "file_path_key":file_path_key}
            data_list.append(item_dict)
            # torchaudio.save(wav_path, wav, 16000)
            with open(wav_path, "wb") as wav_file:
                wav_file.write(wav)
    time_elapsed =dragon_ai.do_get_elapsed_time(time_now)
    dragon_ai.print_info(f'time_elapsed: {time_elapsed}')  # 10个tar包用时13.5s
    dragon_ai.write_dict_list_to_jsonl(data_list, f'{output_tar_dir}/data.list')

def get_old_data_shards_v2():
    """"""
    data_list_no_wav = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_thinking/doubao/data_old/data.list"
    dict_list = dragon_ai.load_dict_list_from_jsonl(data_list_no_wav)
    new_dict_list_path = "/home/A02_tmpdata3/osum_s2s/gender_old/data.list"
    new_dict_list = []
    for dict_i in dict_list:
        key = dict_i['key']
        wav_path = f'/home/A02_tmpdata3/osum_s2s/gender_old/wav/{key}.wav'
        dict_i['wav'] = wav_path
        new_dict_list.append(dict_i)
    dragon_ai.write_dict_list_to_jsonl(new_dict_list, new_dict_list_path)

def get_old_data_shards_v3():
    """"""
    data_list_path = "/home/A02_tmpdata3/osum_s2s/gender_old/data.list"
    # dragon_ai.do_make_shards_common(
    #     data_list_path,
    #     "/home/A02_tmpdata3/osum_s2s/gender_old",
    #     num_utts_per_shard=1000,
    #     prefix='shards',
    #     resample=16000,
    #     num_threads=8,
    #     do_recover=True
    # )
    shards_path = "/home/A02_tmpdata3/osum_s2s/gender_old/shards_list.txt"
    dragon_ai.do_convert_shards_to_format_raw(shards_path, "/home/A02_tmpdata3/osum_s2s/gender_old", batch_size=50, num_workers=10)
    # do fenpei
    task_tag = '<S2TCHAT> <THINKER>'
    output_dir = "/home/A02_tmpdata3/osum_s2s/gender_old/xlgeng_new_data/s2t_thinking"
    tar_dir = "/home/A02_tmpdata3/osum_s2s/gender_old"
    dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, output_dir, task_tag=task_tag, tar_dir=tar_dir)
    task_tag = '<S2TCHAT>'
    output_dir = "/home/A02_tmpdata3/osum_s2s/gender_old/xlgeng_new_data/s2t_no_thinking"
    dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, output_dir, task_tag=task_tag, tar_dir=tar_dir)
    task_tag = "<S2TCHAT> <TEXT2TOKEN> <THINK>"
    output_dir = "/home/A02_tmpdata3/osum_s2s/gender_old/xlgeng_new_data/s2s_think"
    dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, output_dir, task_tag=task_tag, tar_dir=tar_dir)
    task_tag = "<S2TCHAT> <TEXT2TOKEN>"
    output_dir = "/home/A02_tmpdata3/osum_s2s/gender_old/xlgeng_new_data/s2s_no_think"
    dragon_ai.do_distribute_datalist_for_conbine_type(data_list_path, output_dir, task_tag=task_tag, tar_dir=tar_dir)

def do_get_a_few_child_data_raw():
    """"""
    data_list_path = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_no_thinking/doubao/data_child/data_few.list"
    combines_path = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_no_thinking/doubao/data_child/combines_list.txt"
    combines_list = dragon_ai.load_list_from_file_clean(combines_path)
    tar_list = []
    for combine_path in combines_list:
        tar_list.append(os.path.join("/home/A02_tmpdata3/osum_s2s/gender", dragon_ai.do_get_file_pure_name_from_path(combine_path)+".tar"))
    fake_tar_path = dragon_ai.do_get_fake_file_path()
    dragon_ai.write_list_to_file(tar_list, fake_tar_path)
    test_data_loader = dragon_ai.do_get_data_loader_for_read(fake_tar_path, batch_size=50, num_workers=10)
    dict_info_list = dragon_ai.load_dict_list_from_jsonl(data_list_path)
    wav_dir = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_no_thinking/doubao/data_child/wav"
    dragon_ai.makedirs(wav_dir)
    big_dict = {}

    for dict_info in dict_info_list:
        big_dict[dict_info['key']] = dict_info
    keys_set = set(big_dict.keys())
    for i, batch in enumerate(test_data_loader):
        sorted_key = batch['keys']
        sorted_txt = batch['labels']
        sorted_wav = batch['wavs']
        sorted_extra_dicts = batch['extra_dicts']
        sorted_file_path = batch['tar_files']
        j = 0
        for key, txt, wav ,extra_dict,file_path in zip(sorted_key, sorted_txt, sorted_wav,sorted_extra_dicts,sorted_file_path):
            if key not in keys_set:
                continue
            wav_path = os.path.join(wav_dir, f'{key}.wav')
            # 写入wav文件
            with open(wav_path, "wb") as wav_file:
                wav_file.write(wav)
            # 更新字典
            big_dict[key]['wav'] = wav_path
            j += 1
    new_dict_list = list(big_dict.values())
    new_dict_list_path = "/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2s_no_thinking/doubao/data_child/data_few_wav.list"
    dragon_ai.write_dict_list_to_jsonl(new_dict_list, new_dict_list_path)


def do_get_datalist_for_test_data():
    text_scp = "/home/A02_tmpdata2/data/test_record_xlgeng/text.scp"
    wav_scp = "/home/A02_tmpdata2/data/test_record_xlgeng/wav.scp"
    wav_list = dragon_ai.load_list_from_file_clean(wav_scp)
    wav_dict = {}
    for wav_path in wav_list:
        key = dragon_ai.do_get_file_pure_name_from_path(wav_path)
        wav_dict[key] = wav_path
    dragon_ai.write_dict_to_scp(wav_dict, wav_scp)

    text_dict = dragon_ai.load_dict_from_scp(text_scp)
    wav_dict = dragon_ai.load_dict_from_scp(wav_scp)
    data_list = []
    for key in text_dict:
        txt = text_dict[key]
        wav_path = wav_dict[key]
        item_dict = {'key': key, 'txt': txt, 'wav': wav_path}
        data_list.append(item_dict)
    data_list_path = "/home/A02_tmpdata2/data/test_record_xlgeng/data.list"
    dragon_ai.write_dict_list_to_jsonl(data_list, data_list_path)
if __name__ == '__main__':
    do_get_datalist_for_test_data()